package com.xuan.chatroom.interceptor;

import com.xuan.chatroom.constants.Constant;
import com.xuan.chatroom.model.User;
import com.xuan.chatroom.utils.JwtUtil;
import io.jsonwebtoken.Claims;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;

@Component
public class SaveTokenInterceptor implements HandshakeInterceptor {
    // 握手前的操作
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 1. 从 header中获取子协议的token，若token为空、过期或非法的，则拒绝建立 WebSocket连接
        String token = request.getHeaders().getFirst("Sec-WebSocket-Protocol");
        Claims claims = JwtUtil.parseToken(token);
        if (claims == null) return false;

        // 2. 将 id 和 username 存入WebSocket的 attributes中
//        int id = (int) claims.get(Constant.CLAIM_USERID);
//        String username = (String) claims.get(Constant.CLAIM_USERNAME);

        // 2. 将 token 和 id 存入WebSocket的 attributes中
        attributes.put(Constant.WEBSOCKET_USER_TOKEN, token);
        // 额外存储 id 的目的是当用户携带的令牌过期时，服务器主动关闭连接后，能直接利用 id 移除哈希表中的 WebSocketSession，因为令牌过期也就得不到里面的信息
        attributes.put(Constant.WEBSOCKET_USER_ID, claims.get(Constant.CLAIM_USERID));
        return true;
    }


    // 握手完成后的操作
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        // 若前端使用子协议传递 token，则后端也必须子协议作为响应的 header 返回，否则会出现连接异常！！！！！！！！！！！！！！！！
        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
        if (httpRequest.getHeader("Sec-WebSocket-Protocol") != null) {
            httpResponse.addHeader("Sec-WebSocket-Protocol", httpRequest.getHeader("Sec-WebSocket-Protocol"));
        }
    }
}
